import onnxruntime as ort

# 打印 onnxruntime 版本（适应不同版本的方式）
try:
    version = ort.__version__
except AttributeError:
    try:
        version = ort.get_version()
    except AttributeError:
        version = "未知 (无法获取版本信息)"

print(f"ONNXRuntime 版本: {version}")

# 获取可用的执行提供程序列表
providers = ort.get_available_providers()
print(f"可用执行提供程序: {providers}")

# 检查是否支持 CUDA
if 'CUDAExecutionProvider' in providers:
    print("✅ GPU 可用于 ONNX Runtime 推理")
    
    # 创建一个会话，明确指定使用 CUDA
    try:
        # 注意：这里需要一个 .onnx 模型文件才能完全测试 GPU 是否工作
        # 对于测试目的，我们尝试获取 CUDA 设备信息而不创建会话
        cuda_provider_options = {
            "device_id": 0,  # 使用第一个GPU
        }
        
        print("CUDA 配置: ", cuda_provider_options)
    except Exception as e:
        print(f"配置 CUDA 选项时出错: {e}")
        
    # 获取 GPU 设备信息
    import subprocess
    try:
        result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE)
        print("\nNVIDIA GPU 信息:")
        print(result.stdout.decode('utf-8'))
    except Exception as e:
        print(f"无法获取 GPU 信息: {e}")
else:
    print("❌ GPU 不可用于 ONNX Runtime 推理")
    print("仅可用的提供程序:", providers)

